/**
 * Copyright 2023-2023 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph/op/image_para_utils.h"
#include "tensor/image_process_config_types.h"
#include "graph/tensor.h"

#include "infra/base/assertion.h"
#include "securec.h"

namespace hiai {
template<typename T>
    static Buffer CreateUnifiedPara(T* paraData, const size_t paraNums, const ImageFormat& imageFormat);

template<typename T>
Buffer CreateUnifiedPara(T* paraData, const size_t paraNums, const ImageFormat& imageFormat)
{
    size_t tensorSize = paraNums * sizeof(T);
    tensorSize += sizeof(ImageFormat);

    Buffer tensorBuffer = Buffer(tensorSize);
    std::uint8_t* tensorData = tensorBuffer.MutableData();

    HIAI_EXPECT_TRUE_R(
        memcpy_s(tensorData, sizeof(ImageFormat), &imageFormat, sizeof(ImageFormat)) == 0, tensorBuffer);

    size_t memOffset = sizeof(ImageFormat);
    HIAI_EXPECT_TRUE_R(memcpy_s(
        tensorData + memOffset, paraNums * sizeof(T), paraData, paraNums * sizeof(T)) == 0, tensorBuffer);

    return tensorBuffer;
}

namespace {
hiai::op::Const CreateConstFromBuffer(const Buffer& tensorBuffer, const string& name)
{
    TensorDesc tensorDesc(Shape({1, static_cast<int64_t>(tensorBuffer.size()), 1, 1}), FORMAT_NCHW, DT_UINT8);

    TensorPtr outputTensorPtr = std::make_shared<Tensor>(tensorDesc, tensorBuffer);
    hiai::op::Const constData = hiai::op::Const(name).set_attr_value(outputTensorPtr);

    return constData;
}

hiai::op::ConfigData CreateConfigFromBuffer(const Buffer& tensorBuffer, const string& name, const string& type)
{
    TensorDesc tensorDesc(Shape({1, static_cast<int64_t>(tensorBuffer.size()), 1, 1}), FORMAT_NCHW, DT_UINT8);

    TensorPtr outputTensorPtr = std::make_shared<Tensor>(tensorDesc, tensorBuffer);
    hiai::op::ConfigData configdata = hiai::op::ConfigData(name).set_attr_value(outputTensorPtr).set_attr_type(type);

    return configdata;
}
}

hiai::op::Const ImageParaUtil::CreateConfigConst(MultiCropPara& para, const string& name)
{
    Buffer tensorBuffer = CreateUnifiedPara<CropPara>(para.cropParas.data(), para.cropParas.size(), para.imageFormat);
    return CreateConstFromBuffer(tensorBuffer, name);
}

hiai::op::Const ImageParaUtil::CreateConfigConst(MultiResizePara& para, const string& name)
{
    Buffer tensorBuffer =
        CreateUnifiedPara<ResizePara>(para.resizeParas.data(), para.resizeParas.size(), para.imageFormat);
    return CreateConstFromBuffer(tensorBuffer, name);
}

hiai::op::Const ImageParaUtil::CreateConfigConst(MultiDtcPara& para, const string& name)
{
    Buffer tensorBuffer = CreateUnifiedPara<DtcPara>(para.dtcParas.data(), para.dtcParas.size(), para.imageFormat);
    return CreateConstFromBuffer(tensorBuffer, name);
}

hiai::op::Const ImageParaUtil::CreateConfigConst(MultiRotatePara& para, const string& name)
{
    Buffer tensorBuffer =
        CreateUnifiedPara<RotatePara>(para.rotateParas.data(), para.rotateParas.size(), para.imageFormat);
    return CreateConstFromBuffer(tensorBuffer, name);
}

hiai::op::Const ImageParaUtil::CreateConfigConst(MultiPadPara& para, const string& name)
{
    Buffer tensorBuffer = CreateUnifiedPara<PadPara>(para.padParas.data(), para.padParas.size(), para.imageFormat);
    return CreateConstFromBuffer(tensorBuffer, name);
}

hiai::op::Const ImageParaUtil::CreateConfigConst(CscPara& para, const string& name)
{
    TensorDesc tensorDesc(Shape({1, sizeof(CscPara), 1, 1}), FORMAT_NCHW, DT_UINT8);
    TensorPtr outputTensorPtr =
        std::make_shared<Tensor>(tensorDesc, reinterpret_cast<uint8_t*>(&para), sizeof(CscPara));
    hiai::op::Const constData = hiai::op::Const(name).set_attr_value(outputTensorPtr);

    return constData;
}

hiai::op::Const ImageParaUtil::CreateConfigConst(ChannelSwapPara& para, const string& name)
{
    TensorDesc tensorDesc(Shape({1, sizeof(ChannelSwapPara), 1, 1}), FORMAT_NCHW, DT_UINT8);
    TensorPtr outputTensorPtr =
        std::make_shared<Tensor>(tensorDesc, reinterpret_cast<uint8_t*>(&para), sizeof(ChannelSwapPara));
    hiai::op::Const constData = hiai::op::Const(name).set_attr_value(outputTensorPtr);

    return constData;
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(MultiCropPara& para, const string& name, const string& type)
{
    Buffer tensorBuffer = CreateUnifiedPara<CropPara>(para.cropParas.data(), para.cropParas.size(), para.imageFormat);
    return CreateConfigFromBuffer(tensorBuffer, name, type);
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(MultiResizePara& para, const string& name, const string& type)
{
    Buffer tensorBuffer =
        CreateUnifiedPara<ResizePara>(para.resizeParas.data(), para.resizeParas.size(), para.imageFormat);
    return CreateConfigFromBuffer(tensorBuffer, name, type);
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(MultiDtcPara& para, const string& name, const string& type)
{
    Buffer tensorBuffer = CreateUnifiedPara<DtcPara>(para.dtcParas.data(), para.dtcParas.size(), para.imageFormat);
    return CreateConfigFromBuffer(tensorBuffer, name, type);
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(MultiRotatePara& para, const string& name, const string& type)
{
    Buffer tensorBuffer =
        CreateUnifiedPara<RotatePara>(para.rotateParas.data(), para.rotateParas.size(), para.imageFormat);
    return CreateConfigFromBuffer(tensorBuffer, name, type);
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(MultiPadPara& para, const string& name, const string& type)
{
    Buffer tensorBuffer = CreateUnifiedPara<PadPara>(para.padParas.data(), para.padParas.size(), para.imageFormat);
    return CreateConfigFromBuffer(tensorBuffer, name, type);
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(CscPara& para, const string& name, const string& type)
{
    TensorDesc tensorDesc(Shape({1, sizeof(CscPara), 1, 1}), FORMAT_NCHW, DT_UINT8);
    TensorPtr outputTensorPtr =
        std::make_shared<Tensor>(tensorDesc, reinterpret_cast<uint8_t*>(&para), sizeof(CscPara));
    hiai::op::ConfigData configdata = hiai::op::ConfigData(name).set_attr_value(outputTensorPtr).set_attr_type(type);

    return configdata;
}

hiai::op::ConfigData ImageParaUtil::CreateConfigData(ChannelSwapPara& para, const string& name, const string& type)
{
    TensorDesc tensorDesc(Shape({1, sizeof(ChannelSwapPara), 1, 1}), FORMAT_NCHW, DT_UINT8);
    TensorPtr outputTensorPtr =
        std::make_shared<Tensor>(tensorDesc, reinterpret_cast<uint8_t*>(&para), sizeof(ChannelSwapPara));
    hiai::op::ConfigData configdata = hiai::op::ConfigData(name).set_attr_value(outputTensorPtr).set_attr_type(type);

    return configdata;
}

} // namespace hiai